#!/usr/bin/env python
#
# testing for keep-alive connection
# test for gate-client-dispatch model
#

import time,socket,base64,fcntl,errno
import threading
import Queue as queue

import celtuce.net.ioloop as ioloop
import celtuce.net.iostream as iostream

from celtuce.help import logging
logging.getLogger().setLevel(logging.LEVEL_DEBUG)


URL = 'tcp://127.0.0.1:14300'

#######################################################################


class   ServerThread(threading.Thread) :
    def __init__(self, to_agent_queue, tcp_port=0) :
        threading.Thread.__init__(self) 

        self._tcp_server = None

        self._io_loop = ioloop.IOLoop()

        if tcp_port :
            self._tcp_port = tcp_port
            self._tcp_server = TcpServer(to_agent_queue, self._io_loop)

        self._logger = logging.getLogger('serverthread')
        ## __init__()

    def run(self) :
        self._tcp_server.listen( self._tcp_port )

        self._io_loop.start()
        self._logger.info('IOLoop stop.')
        ## run()

    def stop(self) :
        self._logger.info('stop IOLoop ...')
        self._io_loop.stop()
        ## stop()

    def handle(self, response) :
        self._tcp_server.handle( response )
        ## handle()

    ## class ServerThread


class   TcpServer(object) :
    def __init__(self, to_agent_queue, io_loop) :
        self._to_agent_queue = to_agent_queue
        self._io_loop = io_loop

        self._socket = None
        self._connections = {}

        self._logger = logging.getLogger('tcpserver')
        ## __init__()

    def listen(self, port, address='') :
        assert not self._socket
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        flags = fcntl.fcntl(self._socket.fileno(), fcntl.F_GETFD)
        flags |= fcntl.FD_CLOEXEC
        fcntl.fcntl(self._socket.fileno(), fcntl.F_SETFD, flags)
        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._socket.setblocking(0)
        self._socket.bind((address, port))
        self._socket.listen(128)
        self._io_loop.add_handler(self._socket.fileno(), self._handleEvents, self._io_loop.READ)
        ## listen()

    def _handleEvents(self, fd, events):
        while True:
            try:
                connection, address = self._socket.accept()
            except socket.error, e:
                if e[0] in (errno.EWOULDBLOCK, errno.EAGAIN):
                    return
                raise
            try:
                stream = iostream.IOStream(connection, io_loop=self._io_loop)
                self._connections[ address ] = TcpConnection(stream, address, self._to_agent_queue, self.onConnectionClosed)
            except:
                self._logger.excpt()
        ## _handleEvents()

    def handle(self, response) :
        address,data = response
        connection = self._connections.get(address,None)
        if not connection :
            return
        connection.handleResponse(data)
        ## handle()

    def onConnectionClosed(self, address) :
        if address in self._connections :
            del self._connections[ address ]
        ## onConnectionClosed()

    ## class TcpServer

class   TcpConnection(object) :
    def __init__(self, stream, address, to_agent_queue, closed_callback=None) :
        self._stream = stream
        self._address = address
        self._to_agent_queue = to_agent_queue
        self._closed_callback = closed_callback

        self._END_TAG = "\r\n\r\n"
        self._request_finished = False
        self._stream.read_until(self._END_TAG, self.__onRequest)

        self._keep_alive = True

        #self._sending_lock = threading.Lock()

        self._logger = logging.getLogger('tcpconnect')
        self._logger.debug('create TcpConnection for %s', self._address)
        ## __init__()

    def __del__(self) :
        if self._stream.closed() :
            self._logger.debug('stream is closed when destroying')
        else :
            self._logger.debug('stream is opened when destroying')
        self._logger.debug('destroy TcpConnection for %s', self._address)
        pass
        ## __del__()

    def __onRequest(self, data) :
        request = data
        #self._logger.debug('recv request  : %s', repr(request))
        self._to_agent_queue.put((self._address, request.strip()))
        ## __onRequest()

    def handleResponse(self, response) :
        #self._sending_lock.acquire()
        self.write(response)
        ## handleResponse()

    def write(self, data) :
        if data :
            response = data 
            self._logger.debug('ID=%s send response : %s', id(self._stream),repr(response))
            self._stream.write(''.join((response, self._END_TAG)), self.__onWriteComplete)
        else :
            #self._sending_lock.release()
            self._logger.debug('WHY send nothing : %s', repr(data))
            if not self._keep_alive :
                self.finish()
        ## write()

    def finish(self):
        if self._stream and not self._stream.writing():
            self.__finishRequest()

    def __onWriteComplete(self):
        #self._logger.debug('write complete.')
        #self._sending_lock.release()
        if self._keep_alive :
            self._logger.debug('%s waiting next request ...', self._address)
            self._request_finished = False
            self._stream.read_until(self._END_TAG, self.__onRequest)
        else :
            self.__finishRequest()
        ## __onWriteComplete()

    def __finishRequest(self) :
        #self._logger.debug('finish request ...')
        self._stream.close()
        self._stream = None
        if self._closed_callback :
            self._closed_callback( self._address )
        ## __finishRequest()

    ## class TcpConnection


#######################################################################

class   AgentThread(threading.Thread) :
    def __init__(self, response_queue) :
        threading.Thread.__init__(self) 

        self._response_queue = response_queue

        self._io_loop = ioloop.IOLoop()

        self._logger = logging.getLogger('agentthread')
        ## __init__()

    def run(self) :
        self._io_loop.start()
        self._logger.info('IOLoop stop.')
        ## run()

    def stop(self) :
        self._logger.info('stop IOLoop ...')
        self._io_loop.stop()
        ## stop()

    def handle(self, request) :
        # just return
        self._logger.debug('handle request : %s', repr(request))
        self._response_queue.put( request )
        ## handle()

    ## class AgentThread


#######################################################################

class   Dispatcher(threading.Thread) :
    def __init__(self, request_queue, agent) :
        threading.Thread.__init__(self)

        self._request_queue = request_queue
        self._agent = agent

        self._poll_timeout = 1.0
        self._will_stop = False
        ## __init__()

    def stop(self) :
        self._will_stop = True
        ## stop()

    def run(self) :
        while not self._will_stop :
            try :
                msg = self._request_queue.get(True, self._poll_timeout) 
            except queue.Empty :
                continue

            self._agent.handle(msg)
            self._request_queue.task_done()

        # will stop, and discard all msg in queue
        while True :
            try :
                msg = self._request_queue.get(False)
            except queue.Empty :
                break
        ## run()

    ## class RequestDispatcher

#######################################################################

class   Client(object) :
    def __init__(self, server_host, server_port) :
        self._server_host = server_host
        self._server_port = server_port

        self._END_TAG = "\r\n\r\n"

        self._socket = None
        self._keep_alive = True

        self._logger = logging.getLogger('client')
        ## __init__()

    def __del__(self) :
        self.__disconnect(force=True)
        ## __del__()

    ####-----------------------------------------------------------####

    def __connect(self) :
        if self._socket : 
            #self._stream = iostream.IOStream( self._socket, self._io_loop )
            return

        try :
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            s.connect( (self._server_host, self._server_port) )
        except :
            self._logger.excpt()
            raise
        self._socket = s

        self._io_loop = ioloop.IOLoop()
        self._stream = iostream.IOStream( self._socket, self._io_loop )

        return
        ## __connect()

    def __disconnect(self, force=False) :
        if self._keep_alive and not force :
            return

        if self._stream :
            self._stream.close()
            self._stream = None

        if self._socket :
            try :
                self._socket.close()
            except :
                self._logger.excpt()

        self._socket = None
        self._io_loop = None
        ## __disconnect()

    def rpc(self, request) :
        self.__connect()

        self._logger.debug('write request : %s', request)
        self._stream.write( ''.join((request, self._END_TAG)), self.__onWriteComplete )
        self._stream.read_until(self._END_TAG, self.__onResponse)
        
        self._io_loop.start()

        self.__disconnect()
        return self._response
        ## rpc()

    def __onWriteComplete(self) :
        #self._logger.debug('write done.')
        pass
        ## __onWriteComplete()

    def __onResponse(self, response_str) :
        self._logger.debug('read response : %s', repr(response_str))
        response = response_str.strip()
        if response :
            self._response = response_str.strip()
            self._io_loop.stop()
        else :
            self._stream.read_until(self._END_TAG, self.__onResponse)
        ## __onResponse()

    ## class Client


#######################################################################


def echo(any_data) :
    return any_data.strip()
    ## echo()


TCP_PORT = 24344

def testDispatcherSpeed(tranx, from_queue, to_queue) :
    time_start = time.time()
    times = 100000

    request = 'hello-world'
    for i in xrange(times) :
        from_queue.put( request )
        response = to_queue.get()

    time_last = time.time() - time_start
    tps = times / time_last
    print 'dispatcher-speed tps=%.2f' % (tps,)
    ## testDispatcherSpeed()

def runTest() :
    logger = logging.getLogger('test')

    to_agent_queue = queue.Queue()
    to_client_queue = queue.Queue()

    logger.info('test keep-alive connection startup ...')
    server = ServerThread(to_agent_queue, TCP_PORT)
    server.start()

    agent = AgentThread(to_client_queue)
    agent.start()

    request_dispatcher = Dispatcher(to_agent_queue, agent)
    request_dispatcher.start()

    #testDispatcherSpeed(request_dispatcher, to_agent_queue, to_client_queue)

    response_dispatcher = Dispatcher(to_client_queue, server)
    response_dispatcher.start()

    time.sleep(0.5)

    client = Client('', TCP_PORT)

    try :
        # ping
        request = 'hello,world'
        response = client.rpc( request )
        logger.info('REQUEST  : %s', request)
        logger.info('RESPONSE : %s', response)

        # test tps
        request_prefix = 'hello,world'
        time_start = time.time()
        times = 10000
        sum_bytes = 0
        for i in range(times) :
            request = '-'.join((request_prefix, str(i)))
            response = client.rpc( request )
            assert request == response
            sum_bytes += len(response)
        time_last = time.time() - time_start

        tps = times / time_last
        bw = sum_bytes / time_last / 1024.0
        logger.info('rpc %d times last %.2f seconds, tps:%.2f, bw:%.2f KB/s',times,time_last,tps,bw)
    finally :
        logger.info('shutdown server ...')
        server.stop()
        logger.info('joining server-thread ...')
        if server.isAlive() :
            server.join()

        logger.info('shutdown agent ...')
        agent.stop()
        logger.info('joining agent-thread ...')
        if agent.isAlive() :
            agent.join()

        logger.info('shutdown request-dispatcher ...')
        request_dispatcher.stop()
        logger.info('joining request-dispatcher-thread ...')
        if request_dispatcher.isAlive() :
            request_dispatcher.join()

        logger.info('shutdown response-dispatcher ...')
        response_dispatcher.stop()
        logger.info('joining response-dispatcher-thread ...')
        if response_dispatcher.isAlive() :
            response_dispatcher.join()

    logger.info('test keep-alive connection done')
    ## runTest()

if  __name__ == '__main__' :
    runTest()

